import cv2
import numpy as np
import os
from os.path import exists
from imutils import paths
import pickle
from tqdm import tqdm
import logging

# 动态导入tensorflow相关模块
try:
    from tensorflow.keras.applications.vgg16 import VGG16, preprocess_input
    from tensorflow.keras.preprocessing import image

    TENSORFLOW_AVAILABLE = True
except ImportError:
    TENSORFLOW_AVAILABLE = False
    logging.warning("TensorFlow未安装，将无法使用VGG特征提取")

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


def get_size(file):
    size_bytes = os.path.getsize(file)
    size_mb = size_bytes / (1024 * 1024)
    return size_mb


def createXY(train_folder, dest_folder, method='vgg', batch_size=64):
    """
    从训练文件夹中读取图像，生成特征向量和标签
    train_folder参数现在指向包含train和test文件夹的父目录
    """
    x_file_path = os.path.join(dest_folder, "X.pkl")
    y_file_path = os.path.join(dest_folder, "y.pkl")

    # 如果 X 和 y 已经存在，则直接读取
    if os.path.exists(x_file_path) and os.path.exists(y_file_path):
        logging.info("X和y已经存在，直接读取")
        logging.info(f"X文件大小:{get_size(x_file_path):.2f}MB")
        logging.info(f"y文件大小:{get_size(y_file_path):.2f}MB")

        with open(x_file_path, 'rb') as f:
            X = pickle.load(f)
        with open(y_file_path, 'rb') as f:
            y = pickle.load(f)

        return X, y

    logging.info("读取所有图像，生成X和y")

    X = []
    y = []

    # 根据传入的方法选择不同的模型
    if method == 'vgg':
        if not TENSORFLOW_AVAILABLE:
            logging.error("TensorFlow未安装，无法使用VGG特征提取，回退到flat方法")
            method = 'flat'
            model = None
        else:
            model = VGG16(weights='imagenet', include_top=False, pooling="max")
            logging.info("完成构建 VGG16 模型")
    elif method == 'flat':
        model = None

    # 检查数据文件夹结构并收集所有图像
    image_paths_with_labels = []

    # train_folder现在应该是项目根目录，包含train和test文件夹
    base_dir = train_folder  # 项目根目录

    # 检查train文件夹（猫图片）
    train_cat_folder = os.path.join(base_dir, "train")
    if os.path.exists(train_cat_folder):
        cat_paths = list(paths.list_images(train_cat_folder))
        for path in cat_paths:
            image_paths_with_labels.append((path, 0))  # 猫标签为0
        logging.info(f"从 {train_cat_folder} 找到 {len(cat_paths)} 张猫图片")
    else:
        logging.warning(f"猫图片文件夹不存在: {train_cat_folder}")

    # 检查test文件夹（狗图片）
    test_dog_folder = os.path.join(base_dir, "test")
    if os.path.exists(test_dog_folder):
        dog_paths = list(paths.list_images(test_dog_folder))
        for path in dog_paths:
            image_paths_with_labels.append((path, 1))  # 狗标签为1
        logging.info(f"从 {test_dog_folder} 找到 {len(dog_paths)} 张狗图片")
    else:
        logging.warning(f"狗图片文件夹不存在: {test_dog_folder}")

    if len(image_paths_with_labels) == 0:
        logging.error("未找到任何图像文件，创建示例数据")
        return create_sample_data(dest_folder, method)

    # 打乱数据
    np.random.shuffle(image_paths_with_labels)
    image_paths = [item[0] for item in image_paths_with_labels]
    labels = [item[1] for item in image_paths_with_labels]

    # 计算需要的批次数
    num_batches = len(image_paths) // batch_size + (1 if len(image_paths) % batch_size else 0)

    # 使用进度条对批次进行循环处理
    for idx in tqdm(range(num_batches), desc="读取图像"):
        batch_images = []
        batch_labels = []

        start = idx * batch_size
        end = min((idx + 1) * batch_size, len(image_paths))

        # 对于每个批次中的图像
        for i in range(start, end):
            image_path, label = image_paths[i], labels[i]
            try:
                if method == 'vgg':
                    img = image.load_img(image_path, target_size=(224, 224))
                    img = image.img_to_array(img)
                elif method == 'flat':
                    img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
                    if img is None:
                        logging.warning(f"无法读取图像: {image_path}")
                        continue
                    img = cv2.resize(img, (32, 32))
                batch_images.append(img)
                batch_labels.append(label)
            except Exception as e:
                logging.warning(f"处理图像 {image_path} 时出错: {e}")
                continue

        if len(batch_images) == 0:
            continue

        batch_images = np.array(batch_images)
        if method == 'vgg':
            batch_images = preprocess_input(batch_images)
            batch_pixels = model.predict(batch_images, verbose=0)
        else:
            batch_pixels = batch_images.reshape(batch_images.shape[0], -1)

        X.extend(batch_pixels)
        y.extend(batch_labels)

    if len(X) == 0:
        logging.error("未能成功读取任何图像数据")
        return create_sample_data(dest_folder, method)

    X = np.array(X)
    y = np.array(y)

    logging.info(f"X.shape: {X.shape}")
    logging.info(f"y.shape: {y.shape}")
    logging.info(f"猫图片数量: {np.sum(y == 0)}, 狗图片数量: {np.sum(y == 1)}")

    # 确保目标文件夹存在
    os.makedirs(dest_folder, exist_ok=True)

    # 将X和y分别序列化到文件
    with open(x_file_path, 'wb') as f:
        pickle.dump(X, f)
    with open(y_file_path, 'wb') as f:
        pickle.dump(y, f)

    logging.info(f"X和y已保存到 {dest_folder}")

    return X, y


def create_sample_data(dest_folder, method='flat'):
    """创建示例数据用于测试"""
    logging.info("创建示例数据...")

    # 创建一些随机数据作为示例
    if method == 'vgg':
        X = np.random.randn(100, 512).astype(np.float32)
    else:
        X = np.random.randn(100, 1024).astype(np.float32)

    # 创建随机标签（0和1）
    y = np.random.randint(0, 2, 100)

    # 保存示例数据
    x_file_path = os.path.join(dest_folder, "X.pkl")
    y_file_path = os.path.join(dest_folder, "y.pkl")

    os.makedirs(dest_folder, exist_ok=True)
    with open(x_file_path, 'wb') as f:
        pickle.dump(X, f)
    with open(y_file_path, 'wb') as f:
        pickle.dump(y, f)

    logging.info(f"创建了示例数据: X.shape={X.shape}, y.shape={y.shape}")
    return X, y